Summary of Changes

Hello @yizhang2077, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a foundational feature for supporting Mamba radix cache (v0) within the SGLang system. The core objective is to enhance the efficiency of KV cache management for models incorporating Mamba architectures. This is achieved by implementing a specialized radix tree that intelligently handles both standard and Mamba-specific KV states, allowing for better resource utilization and faster inference. The changes span across memory allocation, request scheduling, and cache eviction policies, culminating in significant performance gains as evidenced by the provided benchmarks.

Highlights

• Mamba Radix Cache Implementation: Introduced a new MambaRadixCache class to efficiently manage hybrid (full and Mamba) KV cache states, leveraging a radix tree structure for optimized prefix sharing.
• Memory Management Enhancements: Updated memory_pool.py to support Mamba-specific memory allocation, freeing, and state copying/forking using torch.Tensor for improved GPU compatibility and efficiency.
• Scheduler and Policy Integration: Modified the scheduling and batch management logic across schedule_batch.py, schedule_policy.py, and scheduler.py to seamlessly integrate the new MambaRadixCache, including mechanisms for Mamba cache eviction and detailed memory usage tracking.
• Performance Improvements: Benchmarking results demonstrate a notable increase in request throughput (from 2.64 req/s to 3.54 req/s) and input/output token throughput, alongside reduced end-to-end and time-to-first-token latencies when the Mamba radix cache is enabled.
• Unit Testing: Added comprehensive unit tests in test_mamba_unittest.py to validate the functionality of HybridLinearKVPool, MambaPool, and MambaRadixCache, ensuring correctness of allocation, eviction, and prefix matching.

You need to fix the typo and rename the token_msg variables to token_usage_msg:

During testing, I discovered that the server crashed when token_usage reached more than 0.99.

During testing, I discovered that the server crashed when token_usage reached more than 0.99.

Do you have reproduce command? I think token_usage > 0.99 is an abnormal state. (It is too large and other models will crash as well in this state)

During testing, I discovered that the server crashed when token_usage reached more than 0.99.

Do you have reproduce command? I think token_usage > 0.99 is an abnormal state. (It is too large and other models will crash as well in this state)

reproduce command (server H100, tp-size=2):
python bench_serving.py
--host 127.0.0.1
--port 30000
--dataset-name generated-shared-prefix
--gsp-system-prompt-len 1000
--gsp-question-prompt-len 2744
--gsp-output-prompt-len 820
--gsp-num-groups 8
--gsp-prompts-per-group 128

The root cause is incorrect memory availability checking in the Mamba pool. Instead of checking the Mamba pool, only the MHA pool is used, which leads to attempted memory allocation in a full Mamba pool and subsequent server crash due to None being returned from mamba_pool.alloc().

Clode Sonnet's recommendations:

1. Scheduler.check_memory():
   • Check availability of both pools (MHA and Mamba) separately.
2. PrefillAdder.budget_state():
   • For hybrid models with Mamba, check availability of both pools separately.
3. ScheduleBatch.alloc_token_slots():
   • For Mamba pool, use req_to_token_pool.mamba_pool.alloc() instead of token_to_kv_pool_allocator.alloc().
4. MambaRadixCache.match_prefix():
   • Add Mamba pool availability check before allocation.
5. MambaRadixCache.evict_mamba():
   • Add verification that eviction will free sufficient memory.
6. Scheduler._add_request_to_queue():
   • Add Mamba pool availability check before _prefetch_kvcache().
7. HybridReqToTokenPool.alloc():
   • Check Mamba pool availability before allocation.
8. Add logging when Mamba pool memory is insufficient for diagnostics.

@Swipe4057 Mamba pool controls memory capability by setting available size to 3x max_running_requests here . We can control mamba_usage to around 0.66 at most during this benchmark. I have tried your benchmark in Qwen3-Next-80B-A3B-Instruct-FP8 in H100 and it did not crash (but I found another bug)

@Swipe4057 Mamba pool controls memory capability by setting available size to 3x max_running_requests here . We can control mamba_usage to around 0.66 at most during this benchmark. I have tried your benchmark in Qwen3-Next-80B-A3B-Instruct-FP8 in H100 and it did not crash (but I found another bug)

Run the service with the command and try testing again:
environment:
- SGLANG_ENABLE_JIT_DEEPGEMM=1
command:
--model-path /data/models/Qwen3-Next-80B-A3B-Instruct
--served-model-name Qwen3-Next-80B-A3B-Instruct
--cuda-graph-max-bs 512
--cuda-graph-bs 1 2 4 8 16 24 32 40 48 56 64 72 80 88 96 104 112 120 128 136 144 152 160 168 176 184 192 200 208 216 224 232 240 248 256 264 272 280 288 296 304 312 320 328 336 344 352 360 368 376 384 392 400 408 416 424 432 440 448 456 464 472 480 488 496 504 512
--sleep-on-idle
--port 8027
--host 0.0.0.0
--schedule-policy lof
--random-seed 11111
--context-length 131072
--grammar-backend xgrammar
--tool-call-parser qwen25
--enable-metrics
--quantization w8a8_fp8
--allow-auto-truncate
--mamba-ssm-dtype bfloat16
--max-running-requests 1024
--tp-size 2
--ep-size 2
--chunked-prefill-size 16384
--prefill-attention-backend flashinfer
--decode-attention-backend flashinfer
--mem-fraction-static 0.86
--max-running-requests 1024
--api-key 123

Model: https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct

@Swipe4057 I have tried and it is ok. Could you share your server log? (error and mamba_usage in log are important items)

seems like mamba tree cache sanity check is not running, let's add it?

def check_tree_cache(self):
        if self.is_hybrid and isinstance(self.tree_cache, SWARadixCache):
            self.tree_cache.sanity_check()

open sanity_check

@yizhang2077 Could you resolve the conflicts? Then we can merge it.

yizhang2077 hanming-lu ispobock I retested your new code.,
Here's the command to start the server on 2xH100:

--model-path /data/models/Qwen3-Next-80B-A3B-Instruct --served-model-name Qwen3-Next-80B-A3B-Instruct --cuda-graph-max-bs 512 --cuda-graph-bs 1 2 4 8 16 24 32 40 48 56 64 72 80 88 96 104 112 120 128 136 144 152 160 168 176 184 192 200 208 216 224 232 240 248 256 264 272 280 288 296 304 312 320 328 336 344 352 360 368 376 384 392 400 408 416 424 432 440 448 456 464 472 480 488 496 504 512 --sleep-on-idle --port 8027 --host 0.0.0.0 --api-key ${SGLANG_API_KEY} --schedule-policy lof --random-seed 11111 --context-length 131072 --grammar-backend xgrammar --prefill-attention-backend flashinfer --decode-attention-backend flashinfer --tool-call-parser qwen25 --enable-metrics --quantization w8a8_fp8 --allow-auto-truncate --tp-size 2 --ep-size 2 --chunked-prefill-size 16384 --max-running-requests 1024 --mem-fraction-static 0.87 --mamba-ssm-dtype bfloat16 --mamba-full-memory-ratio 4

TEST-1:
reproduce command:
python bench_serving.py
--host 127.0.0.1
--port 30000
--dataset-name generated-shared-prefix
--gsp-system-prompt-len 1000
--gsp-question-prompt-len 2744
--gsp-output-prompt-len 820
--gsp-num-groups 8
--gsp-prompts-per-group 128

results:
main, disable radix:
input_throughput: 18518

mr, enable radix:
input_throughput: 16174 - the results are only getting worse !!!

logs:
Server startup is slow and 503 errors are occurring (this does not happen in main)

TEST-2:
reproduce command:
python bench_serving.py
--host 127.0.0.1
--port 30000
--dataset-name generated-shared-prefix
--gsp-system-prompt-len 1000
--gsp-question-prompt-len 0 - pay attention to this line and the results !!!
--gsp-output-prompt-len 1 - pay attention to this line and the results !!!
--gsp-num-groups 8
--gsp-prompts-per-group 128

results:
main, disable radix:
input_throughput: 45264

mr, enable radix:
input_throughput: 39668

As you can see, the last test sends 8 groups of identical prompts with a length of 1000 characters, without a variable part and with only 1 token for generation. Performance with radix cache enabled dropped dramatically!
And there is simply no cache match!

yizhang2077 Try this code to fix null cache hit:

def match_prefix(self, key: RadixKey, **kwargs) -> MatchResult:
    """Find the matching prefix from the radix tree."""
    cow_mamba: bool = kwargs.get("cow_mamba", False)
    req: Req = kwargs.get("req", None)

    if self.disable or len(key) == 0:
        return MatchResult(
            device_indices=torch.empty(
                (0,),
                dtype=torch.int64,
                device=self.device,
            ),
            last_device_node=self.root_node,
            last_host_node=self.root_node,
        )

    # Get the full matched prefix (including tombstone nodes)
    full_value, full_last_node, mamba_last_node = self._match_prefix_helper(key)

    # copy mamba state to req local space if cow is true
    if cow_mamba and mamba_last_node.mamba_value is not None:
        # for reqs without mamba cache
        if req.mamba_pool_idx is None:
            dst_index = self.req_to_token_pool.mamba_pool.alloc(1)
            # try to alloc again, protect mamba_last_node from eviction
            if dst_index is None:
                self.inc_lock_ref(mamba_last_node)
                self.evict_mamba(1)
                dst_index = self.req_to_token_pool.mamba_pool.alloc(1)
                self.dec_lock_ref(mamba_last_node)
                assert dst_index is not None, "Can not alloc mamba cache"
            src_index = mamba_last_node.mamba_value
            self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index)
            req.mamba_pool_idx = dst_index[0]
        else:
            src_index = mamba_last_node.mamba_value
            dst_index = req.mamba_pool_idx.unsqueeze(0)
            self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index)

    # Optimize: use torch.cat only if needed
    if full_value:
        value = torch.cat(full_value) if len(full_value) > 1 else full_value[0]
    else:
        value = torch.empty((0,), dtype=torch.int64, device=self.device)

    return MatchResult(
        device_indices=value,
        last_device_node=full_last_node,
        last_host_node=full_last_node,
    )

def _match_prefix_helper(
    self, key: RadixKey
) -> Tuple[List[torch.Tensor], TreeNode, TreeNode]:
    """
    Mamba prefix matching helper. Returns:
    - full_value: all matched full cache tokens (including through tombstone nodes)
    - full_last_node: last node with full cache match
    - mamba_last_node: last node with mamba_value (for COW)
    """
    node = self.root_node
    child_key = self.get_child_key_fn(key)
    value = []
    mamba_last_node = node  # Track last node with mamba state

    while len(key) > 0 and child_key in node.children.keys():
        child = node.children[child_key]

        prefix_len = self.key_match_fn(child.key, key)
        if prefix_len < len(child.key):
            new_node = self._split_node(child.key, child, prefix_len)
            value.append(new_node.value)
            node = new_node
            # Update mamba_last_node if new_node has mamba state
            if new_node.mamba_value is not None:
                mamba_last_node = new_node
            break
        else:
            value.append(child.value)
            # Update mamba_last_node before moving to child
            if child.mamba_value is not None:
                mamba_last_node = child
            node = child
            key = key[prefix_len:]
            if len(key):
                child_key = self.get_child_key_fn(key)

    # update time for matched nodes - single pass
    self.full_lru_list.reset_node_and_parents_mru(node, self.root_node)
    if mamba_last_node != self.root_node and mamba_last_node.mamba_value is not None:
        self.mamba_lru_list.reset_node_and_parents_mru(mamba_last_node, self.root_node)

    return value, node, mamba_last_node

TEST-1:

results:
main, disable radix:
input_throughput: 18518

mr+THIS FIX, enable radix:
input_throughput: 16174 up to 17993

TEST-2:

results:
main, disable radix:
input_throughput: 45264

mr+THIS FIX, enable radix:
input_throughput: 39668 up to 298052 !!!!!!!

@Swipe4057 let's merge this PR to add mamba radix cache functionality. welcome to make changes to improve performance!

/data/install/backup/sglang_main/python/sglang/srt/mem_cache/memory_pool.py", line 306, in alloc
mid = self.mamba_pool.alloc(1)[0]
TypeError: 'NoneType' object is not subscriptable

on 20 requests,--max-mamba-cache-size 100000